from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.output_parsers import JsonOutputParser
from typing import List

# 定义单选题模型
class SingleQuestion(BaseModel):
    question_single: str = Field(description="单选题的题目")
    option_single: List[str] = Field(description="单选题的选项列表，包含四个字符串选项")
    answer_single: str = Field(description="单选题的正确答案，字符串格式")

# 定义多个单选题模型，即是一个包含多个单选题的数组
class SingleQuestionList(BaseModel):
    questions: List[SingleQuestion] = Field(description="一个单选题的数组，数组元素是多个单选题")

# 创建JsonOutputParser，解析成上述定义的格式
parser_single = JsonOutputParser(pydantic_object=SingleQuestionList)


# 定义多选题模型类
class MultipleQuestion(BaseModel):
    question_multiple: str = Field(description="多项选择题目")
    option_multiple: List[str] = Field(description="多选题选项列表，包含四个选项")
    answer_multiple: List[str] = Field(description="多选题正确选项，包含一个或多个正确选项")

class MultipleQuestionList(BaseModel):
    questions: List[MultipleQuestion] = Field(description="一个多选题的数组，数组元素是多个多选题")

# 创建JsonOutputParser，解析成上述定义的格式
parser_multiple = JsonOutputParser(pydantic_object=MultipleQuestionList)


class JudgeQuestion(BaseModel):
    # 判断题的问题，为什么判断题要单独出一个，因为如果问题是这样：下列哪项不是运行部全厂停电保障措施中提到的‘组织措施’部分的内容?，判断题出不了了。
    question_judge: str = Field(description="judge的判断题题目，以陈述句的口吻给出一个句子。")
    # option_judge: list = Field(description="判断题列表，列表中仅包含两个元素")  #
    answer_judge: str = Field(description="判断题正确答案，True代表正确,False代表错误")  # 如A

class JudgeQuestionList(BaseModel):
    questions: List[JudgeQuestion] = Field(description="一个判断题的数组，数组元素是多个判断题")

parser_judge = JsonOutputParser(pydantic_object=JudgeQuestionList)